from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from loguru import logger
from pydantic import Field, BaseModel

from a0base.base_llm import pop_llm


class Person(BaseModel):
    """
    定义一个新闻结构化的数据模型类

    属性:
        time (str): 新闻发生的时间
        person (str): 新闻涉及的人物
        event (str): 发生的具体事件
    """
    time: str = Field(description="时间")
    person: str = Field(description="人物")
    event: str = Field(description="事件")


# 1-创建JSON输出解析器，用于将model输出解析为Person对象-获取格式化指令，告诉model如何输出符合要求的JSON格式
json_parser = JsonOutputParser(pydantic_object=Person)
json_format_instructions = json_parser.get_format_instructions()

# 2-创建聊天提示模板，包含系统角色设定和用户问题输入
chat_prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个AI助手，你只能输出结构化JSON数据。"),
    ("human", "请生成一个关于{topic}的新闻。{format_instructions}")
])

# 3-使用指定的角色和问题生成具体的提示内容
prompt = chat_prompt.format_messages(topic="小米", format_instructions=json_format_instructions)
logger.info(prompt)

# 调用模型获取回答结果
response = pop_llm.invoke(prompt)
logger.info(f"解析后的结构化结果:\n{response}")

# 打印类型
logger.info(f"结果类型: {type(response)}")
